{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "djqHSONJ20X8",
    "outputId": "c8c161ce-6a9a-49a8-9b92-b062d9da025f"
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell\n",
    "\n",
    "# install NeMo\n",
    "BRANCH = 'v1.0.0'\n!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[nlp]"
   ]
  },
  {
      "cell_type": "code",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "JFWG-jYCfvD7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# If you're not using Colab, you might need to upgrade jupyter notebook to avoid the following error:\n",
        "# 'ImportError: IProgress not found. Please update jupyter and ipywidgets.'\n",
        "\n",
        "! pip install ipywidgets\n",
        "! jupyter nbextension enable --py widgetsnbextension\n",
        "\n",
        "# Please restart the kernel after running this cell"
      ],
      "execution_count": null,
      "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 124
    },
    "colab_type": "code",
    "id": "CmvtH0pxHDQC",
    "outputId": "fdaddea0-88fe-4207-ee7e-5a6a8ca360bb"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import wget\n",
    "from nemo.collections import nlp as nemo_nlp\n",
    "from omegaconf import OmegaConf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "J6AARcXXUEbs"
   },
   "source": [
    "# Language models\n",
    "\n",
    "Natural Language Processing (NLP) field experienced a huge leap in recent years due to the concept of transfer learning enabled through pretrained language models.\n",
    "\n",
    "[BERT](https://arxiv.org/abs/1810.04805), [RoBERTa](https://arxiv.org/abs/1907.11692), [Megatron-LM](https://arxiv.org/abs/1909.08053), and many other proposed language models achieve state-of-the-art results on many NLP tasks, such as:\n",
    "* question answering\n",
    "* sentiment analysis\n",
    "* named entity recognition and many others.\n",
    "\n",
    "In NeMo, most of the NLP models represent a pretrained language model followed by a Token Classification layer or a Sequence Classification layer or a combination of both. By changing the language model, you can improve the performance of your final model on the specific downstream task you are solving.\n",
    "\n",
    "With NeMo you can use either pretrain a BERT model from your data or use a pretrained language model from [HuggingFace transformers](https://github.com/huggingface/transformers) or [Megatron-LM](https://github.com/NVIDIA/Megatron-LM) libraries.\n",
    "\n",
    "Let's take a look at the list of available pretrained language models, note the complete list of HuggingFace model could be found at [https://huggingface.co/models](https://huggingface.co/models):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "zp7F45bgX7SU",
    "outputId": "b87752df-8e50-400b-9e1a-12e05bf175d2"
   },
   "outputs": [],
   "source": [
    "nemo_nlp.modules.get_pretrained_lm_models_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "OdpxEiirX8F-"
   },
   "source": [
    "NLP models for downstream tasks use `get_lm_model` helper function to easily switch between language models from the list above to another:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "5zllsvBojxuJ",
    "outputId": "1c6e9878-2e82-460a-d890-30073bb11994"
   },
   "outputs": [],
   "source": [
    "# use any pretrained model name from the list above\n",
    "nemo_nlp.modules.get_lm_model(pretrained_model_name='distilbert-base-uncased')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jEgEo0aPj3Ws"
   },
   "source": [
    "All NeMo [NLP models](https://github.com/NVIDIA/NeMo/tree/main/examples/nlp) have an associated config file. As an example, let's examine the config file for the Named Entity Recognition (NER) model (more details about the model and the NER task could be found [here](https://github.com/NVIDIA/NeMo/blob/main/tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "colab_type": "code",
    "id": "T1gA8PsJ13MJ",
    "outputId": "326a3673-a814-422a-bde3-e61ba4f20fc4"
   },
   "outputs": [],
   "source": [
    "MODEL_CONFIG = \"token_classification_config.yaml\"\n",
    "\n",
    "# download the model's configuration file \n",
    "if not os.path.exists(MODEL_CONFIG):\n",
    "    print('Downloading config file...')\n",
    "    wget.download(f'https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/nlp/token_classification/conf/' + MODEL_CONFIG)\n",
    "else:\n",
    "    print ('Config file already exists')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "mX3KmWMvSUQw",
    "outputId": "bf9b174f-f3b1-4387-9c66-29170dec48cf"
   },
   "outputs": [],
   "source": [
    "# this line will print the entire config of the model\n",
    "config = OmegaConf.load(MODEL_CONFIG)\n",
    "print(OmegaConf.to_yaml(config))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PKFbFDwzKUZB"
   },
   "source": [
    "For this tutorial, we are interested in the language_model part of the Named Entity Recognition Model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 104
    },
    "colab_type": "code",
    "id": "v6hlcCYyKFiY",
    "outputId": "0fc8a4e4-978b-406f-e9ef-981f527ccedd"
   },
   "outputs": [],
   "source": [
    "print(OmegaConf.to_yaml(config.model.language_model))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xldsDiH9ZY2h"
   },
   "source": [
    "There might be slight differences from one model to another, but most of them have the following important parameters associated with the language model:\n",
    "* `pretrained_model_name` - a name of the pretrained model from either HuggingFace or Megatron-LM libraries, for example, bert-base-uncased or megatron-bert-345m-uncased.\n",
    "* `lm_checkpoint` - a path to the pretrained model checkpoint if, for example, you trained a BERT model with your data\n",
    "* `config_file` -  path to the model configuration file\n",
    "* `config` or `config_dict` - path to the model configuration dictionary\n",
    "\n",
    "To modify the default language model, specify the desired language model name with the `model.language_model.pretrained_model_name` argument, like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5r_1gIXwashY"
   },
   "outputs": [],
   "source": [
    "config.model.language_model.pretrained_model_name = 'roberta-base'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "EVp4zvxPatga"
   },
   "source": [
    "and then start the training as usual (please see [tutorials/nlp](https://github.com/NVIDIA/NeMo/tree/main/tutorials/nlp) for more details about training of a particular model). \n",
    "\n",
    "You can also provide a pretrained language model checkpoint and a configuration file if available.\n",
    "\n",
    "Note, that `pretrained_model_name` is used to set up both Language Model and Tokenizer.\n",
    "\n",
    "All the above holds for both HuggingFace and Megatron-LM pretrained language models. Let's separately examine some specifics of finetuning with Megatron-LM and HuggingFace models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Wi_LBdSA4U6Q"
   },
   "source": [
    "# Downstream tasks with Megatron and BioMegatron Language Models\n",
    "\n",
    "[Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://arxiv.org/abs/1909.08053) is a large, powerful transformer developed by the Applied Deep Learning Research team at NVIDIA. More details could be found at [Megatron-LM github repo](https://github.com/NVIDIA/Megatron-LM).\n",
    "\n",
    "To see the list of available Megatron-LM models in NeMo, run:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 121
    },
    "colab_type": "code",
    "id": "0o-XPMrIQBmm",
    "outputId": "7a7e0419-e9a1-48b3-b6bf-4d2863f69f46"
   },
   "outputs": [],
   "source": [
    "nemo_nlp.modules.get_megatron_lm_models_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "FM_ei7OSLn1X"
   },
   "source": [
    "If you want to use one of the available Megatron-LM models, specify its name with `model.language_model.pretrained_model_name` argument, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RmcQLoouME8k"
   },
   "outputs": [],
   "source": [
    "config.model.language_model.pretrained_model_name = 'megatron-bert-345m-uncased'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mla-BZexMM2u"
   },
   "source": [
    "If you have a different checkpoint or a model configuration file, use these general Megatron-LM model names:\n",
    "* `megatron-bert-uncased` or \n",
    "* `megatron-bert-cased` \n",
    "\n",
    "and provide associated bert_config and bert_checkpoint files, as follows:\n",
    "\n",
    "`model.language_model.pretrained_model_name=megatron-bert-uncased \\\n",
    "model.language_model.lm_checkpoint=<PATH_TO_CHECKPOINT> \\\n",
    "model.language_model.config_file=<PAHT_TO_CONFIG>`\n",
    " \n",
    " or \n",
    " \n",
    "`model.language_model.pretrained_model_name=megatron-bert-cased \\\n",
    "model.language_model.lm_checkpoint=<PATH_TO_CHECKPOINT> \\\n",
    "model.language_model.config_file=<PAHT_TO_CONFIG>`\n",
    "\n",
    "The general Megatron-LM model names are used to download the correct vocabulary file needed to setup the model correctly. Note, the data preprocessing and model training is done in NeMo. Megatron-LM has its own set of training arguments (including tokenizer) that are ignored during finetuning in NeMo. Please see downstream task [config files and training scripts](https://github.com/NVIDIA/NeMo/tree/main/examples/nlp) for all NeMo supported arguments.\n",
    "\n",
    "## Download pretrained model\n",
    "\n",
    "With NeMo, the original and domain-specific Megatron-LM BERT models and model configuration files will be downloaded automatically, but they also could be downloaded with the links below:\n",
    "\n",
    "[Megatron-LM BERT Uncased 345M (~345M parameters): https://ngc.nvidia.com/catalog/models/nvidia:megatron_bert_345m](https://ngc.nvidia.com/catalog/models/nvidia:megatron_bert_345m/files?version=v0.1_uncased)\n",
    "\n",
    "[Megatron-LM BERT Cased 345M (~345M parameters): https://ngc.nvidia.com/catalog/models/nvidia:megatron_bert_345m](https://ngc.nvidia.com/catalog/models/nvidia:megatron_bert_345m/files?version=v0.1_cased)\n",
    "\n",
    "[BioMegatron-LM BERT Cased 345M (~345M parameters): https://ngc.nvidia.com/catalog/models/nvidia:biomegatron345mcased](https://ngc.nvidia.com/catalog/models/nvidia:biomegatron345mcased)\n",
    "\n",
    "[BioMegatron-LM BERT Uncased 345M (~345M parameters)](https://ngc.nvidia.com/catalog/models/nvidia:biomegatron345muncased): https://ngc.nvidia.com/catalog/models/nvidia:biomegatron345muncased"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "caafJmVLmEVD"
   },
   "source": [
    "# Using any HuggingFace Pretrained Model\n",
    "\n",
    "Currently, there are 4 HuggingFace language models that have the most extensive support in [NeMo](https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/nlp/modules/common/huggingface): \n",
    "\n",
    "* BERT\n",
    "* RoBERTa\n",
    "* ALBERT\n",
    "* DistilBERT\n",
    "\n",
    "As was mentioned before, just set `model.language_model.pretrained_model_name` to the desired model name in your config and get_lm_model() will take care of the rest.\n",
    "\n",
    "If you want to use another language model from [https://huggingface.co/models](https://huggingface.co/models), use HuggingFace API directly in NeMo.\n",
    "More details on model training could be found at [tutorials](https://github.com/NVIDIA/NeMo/tree/main/tutorials)."
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "01_Pretrained_Language_Models_for_Downstream_Tasks.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
